import os
import uuid
from collections.abc import Generator
from typing import Any

import markdown2
import requests
from bs4 import BeautifulSoup
from dify_plugin import Tool
from dify_plugin.entities.tool import ToolInvokeMessage


def convert_html_to_pdf(
    html_content: str,
    page_size: str = "A4",
    gotenberg_url: str = "https://demo.gotenberg.dev",
    wait_delay: str = "5s",
) -> str:
    # Create a temporary file for the HTML content
    file_id = str(uuid.uuid4())

    # Clean up HTML content using BeautifulSoup
    soup = BeautifulSoup(html_content, "html.parser")

    # Add base styling
    style_tag = soup.new_tag("style")
    style_tag.string = """
    body {
        font-family: Arial, sans-serif;
        margin: 20px;
        line-height: 1.6;
    }
    h1, h2, h3, h4, h5, h6 {
        color: #333;
    }
    pre, code {
        background-color: #f5f5f5;
        border-radius: 3px;
        padding: 2px 5px;
    }
    """
    soup.head.append(style_tag)

    clean_html = str(soup)

    # Convert to PDF using Gotenberg
    url = f"{gotenberg_url.rstrip('/')}/forms/chromium/convert/html"
    files = {"files": ("index.html", clean_html, "text/html")}

    data = {
        "waitDelay": wait_delay,
    }

    try:
        response = requests.post(url, files=files, data=data)
        response.raise_for_status()

        output_path = f"/tmp/{file_id}.pdf"
        with open(output_path, "wb") as f:
            f.write(response.content)

        return output_path
    except Exception as e:
        raise e


def convert_markdown_to_pdf(
    markdown_content: str,
    page_size: str = "A4",
    gotenberg_url: str = "https://demo.gotenberg.dev",
    wait_delay: str = "5s",
) -> str:
    # Convert Markdown to HTML
    html_content = markdown2.markdown(
        markdown_content,
        extras=[
            "fenced-code-blocks",
            "tables",
            "header-ids",
            "toc",
            "cuddled-lists",
            "code-friendly",
        ],
    )

    # Wrap HTML content in a proper HTML document
    html_document = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <title>Markdown Document</title>
    </head>
    <body>
        {html_content}
    </body>
    </html>
    """

    # Convert HTML to PDF
    return convert_html_to_pdf(html_document, page_size, gotenberg_url, wait_delay)


class Html2pdfTool(Tool):
    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
        content = tool_parameters.get("content", "")
        content_type = tool_parameters.get("content_type", "markdown")
        filename = tool_parameters.get("filename", "document.pdf")
        page_size = tool_parameters.get("page_size", "A4")
        wait_delay = tool_parameters.get("wait_delay", "5s")
        credentials = self.runtime.credentials
        # Get OSS credentials from provider credentials
        gotenberg_url = credentials.get("gotenberg_url", "https://demo.gotenberg.dev")

        if not content:
            yield self.create_error_message("Content cannot be empty")
            return

        try:
            if content_type == "html":
                output_path = convert_html_to_pdf(content, page_size, gotenberg_url, wait_delay)
            else:  # markdown
                output_path = convert_markdown_to_pdf(content, page_size, gotenberg_url, wait_delay)

            if not output_path or not os.path.exists(output_path):
                yield self.create_error_message("Failed to generate PDF")
                return

            with open(output_path, "rb") as f:
                pdf_data = f.read()

            # 确保文件名以 .pdf 结尾
            if not filename.endswith(".pdf"):
                filename += ".pdf"

            yield self.create_text_message(
                f"Successfully converted to PDF {output_path}"
            )
            # Return the PDF file path
            yield self.create_blob_message(
                pdf_data,
                meta={"filename": filename, "mime_type": "application/pdf"},
            )
        except Exception as e:
            yield self.create_error_message(f"Failed to generate PDF: {e}")
            return